"""
@Description :
对于各个算法实现模块的测试
@Author : inY_Yue
@Time : 2022/5/2-14:35
"""

from LinearRegression import LinearRegression
from KNN import KNN
from LogisticRegression import LogisticRegression
from SupportVectorMachine import SupportVectorMachine
from ClassificationAndRegressiontree.CART import CartClassificationTree
import matplotlib.pyplot as plt

import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler


def test_linear_regression():
    # 生成测试数据
    X_train = np.array([[1], [2], [4], [5]])
    y_train = np.array([2, 4, 6, 8])
    # 构建模型
    regr = LinearRegression.LinearRegression()
    # 训练模型
    regr.train(X_train, y_train)
    # 构建测试数据
    X_test = np.array([[0.5], [1.5], [2.5], [3.5], [4.5], [5.5]])
    # 进行预测
    y_pred = regr.predict(X_test)
    # 绘图
    plt.scatter(X_train, y_train)
    plt.plot(X_test, y_pred)
    plt.savefig('./demos_pic/linearregression.png')
    # plt.show()


def test_logistic_regression():
    X_train = np.array([[1, 0], [5, 1], [6, 4], [4, 2], [3, 2]])
    y_train = np.array([0, 1, 1, 0, 0])
    ss = StandardScaler()
    ss.fit(X_train)
    X_train_std = ss.transform(X_train)
    model = LogisticRegression.LogisticRegression(n_iter=700, eta=0.01)
    model.train(X_train_std, y_train)
    X_test = np.array([[1, 0], [5, 1], [6, 4], [4, 2], [3, 2]])
    X_test_std = ss.transform(X_test)
    print(model.predict(X_test_std))
    # TODO:获取分类概率
    # y_pred_proba = model.predict_probas()


def test_knn():
    model = KNN.KDTree(3)
    df = pd.read_excel('./test_datasets/葡萄酒.xlsx')
    X_train = np.array(df[['酒精含量(%)', '苹果酸含量(%)']])
    y_train = np.array(df['分类'])
    # 数据标准化
    ss = StandardScaler()
    ss.fit(X_train)
    X_train_std = ss.transform(X_train)

    model.train(X_train_std, y_train)
    X_test = [[7, 1], [8, 3]]
    X_test_std = ss.transform(X_test)
    print(model.predict(X_test_std))


def test_svm():
    model = SupportVectorMachine.SMO(C=1, tol=0.01, kernel='rbf', gamma=0.01)
    df = pd.read_excel('./test_datasets/葡萄酒.xlsx')
    X_train = np.array(df[['酒精含量(%)', '苹果酸含量(%)']])
    y_train = np.array(df['分类'])
    ss = StandardScaler()
    ss.fit(X_train)
    X_train_std = ss.transform(X_train)
    model.train(X_train_std, y_train)
    X_test = [[7, 1], [8, 3]]
    X_test_std = ss.transform(X_test)
    print(model.predict(X_test_std))

    x = X_train_std[:, 0]
    y = X_train_std[:, 1]
    plt.scatter(x, y)
    plt.savefig('./demos_pic/svm.png')
    # plt.show()


def test_CART():
    model = CartClassificationTree()
    df = pd.read_excel('./test_datasets/葡萄酒.xlsx')
    X_train = np.array(df[['酒精含量(%)', '苹果酸含量(%)']])
    y_train = np.array(df['分类'])

    ss = StandardScaler()
    ss.fit(X_train)
    X_train_std = ss.transform(X_train)

    model.train(X_train_std, y_train)

    X_test = [[7, 1], [8, 3]]
    X_test_std = ss.transform(X_test)

    print(model.predict(X_test_std))

    x = X_train_std[:, 0]
    y = X_train_std[:, 1]
    plt.scatter(x, y)
    plt.savefig('./demos_pic/cart.png')
    plt.show()


if __name__ == '__main__':
    # test_linear_regression()
    # test_logistic_regression()
    # test_knn()
    # test_svm()
    test_CART()
